from torch import nn
from src.model.attention import Attention


class FeatureBlock(nn.Module):
    def __init__(self, dims):
        super(FeatureBlock, self).__init__()
        self.input_size, self.head, self.output_size = dims
        self.attention = Attention(self.head, self.input_size // self.head)
        self.layer_norm1 = nn.LayerNorm(self.input_size)
        self.fc = nn.Sequential(nn.Linear(self.input_size, self.input_size), nn.ReLU())
        self.layer_norm2 = nn.LayerNorm(self.input_size)

    def forward(self, x, is_training=True):
        B, T, N, C = x.shape
        x = x.permute(0, 2, 3, 1).contiguous()
        x = x.reshape(-1, C, T)
        tx, a = self.attention(x, x, x)
        y = tx + x
        y = self.layer_norm1(y)
        y = y + self.fc(y)
        y = self.layer_norm2(y)

        y = y.reshape(B, N, C, T)
        y = y.permute(0, 3, 1, 2).contiguous()
        return y, a
